/*
 * ID3
 * 
 * Project: Search Engine
 * 
 * Copy right 2012 by Phan Hoàn Vũ
 * Email hoanvuphan@gmail.com
 */

package hcmus.fit.ml.tr.classifier;

import hcmus.fit.ml.tr.data.Attribute;
import hcmus.fit.ml.tr.data.DataSource;
import hcmus.fit.ml.tr.data.Instance;
import hcmus.fit.ml.tr.data.Instances;

import java.util.ArrayList;
import java.util.Calendar;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;

import org.apache.commons.lang3.ObjectUtils;

/**
 * @author hvp
 * 
 */
public class ID3 implements Classifier {

    public class Node {
        private Attribute attribute;
        private Map<Object, Node> branchs;
        private Object decideValue;

        public Node(Attribute attribute, Map<Object, Node> branchs) {
            super();
            this.setAttribute(attribute);
            this.setBranchs(branchs);
        }

        public Attribute getAttribute() {
            return attribute;
        }

        public void setAttribute(Attribute attribute) {
            this.attribute = attribute;
        }

        public Map<Object, Node> getBranchs() {
            return branchs;
        }

        public void setBranchs(Map<Object, Node> branchs) {
            this.branchs = branchs;
        }

        public Object getDecideValue() {
            return decideValue;
        }

        public void setDecideValue(Object decideValue) {
            this.decideValue = decideValue;
        }

        public String toString(int level) {
            StringBuilder b = new StringBuilder();
            if (decideValue != null) {
                b.append(" : ").append(decideValue).append("\n");
            } else if (attribute != null) {
                b.append("\n");
                for (Entry<Object, Node> e : branchs.entrySet()) {
                    for (int i = 0; i < level; i++) {
                        b.append("|\t");
                    }
                    b.append(attribute.getName()).append(" = ").append(e.getKey());
                    b.append(e.getValue().toString( level + 1));
                }
            }
            return b.toString();
        }
        @Override
        public String toString() {
            StringBuilder b = new StringBuilder();
            if (decideValue != null) {
                b.append(decideValue).append("\n");
            } else if (attribute != null) {
                for (Entry<Object, Node> e : branchs.entrySet()) {
                    b.append("|").append(attribute.getName()).append("=").append(e.getKey()).append("\n");
                    b.append(e.getValue());
                }
            }
            return b.toString();
        }
    }

    protected static double log2 = Math.log(2);
    private Node root;

    /**
     * {@inheritDoc}
     */
    public void buildClassifer(Instances data) {

        // 1. Get best attribute
        Attribute bestAttribute = getBestAttribute(data);

        // 2. Tao nut moi. Cap nhat cay
        root = new Node(bestAttribute, new HashMap<Object, Node>());

        Object decideValue = tryGetDecideValue(data);
        if (decideValue != null) {
            // is leaf
            root.setDecideValue(decideValue);
            return;
        }

        // 3. Tao dataSet con
        Map<Object, Instances> childDataSets = buildChildDataSetsBaseOnBestAttribute(data,
                bestAttribute);

        // 4. process child dataSet
        for (Entry<Object, Instances> e : childDataSets.entrySet()) {
            buildTree(e.getValue(), e.getKey(), root);
        }
    }

    private Object tryGetDecideValue(Instances data) {
        if (data.getInstances().size() == 0) {
            return null;
        }
        String nodeName = data.getClassAttribute().getName();
        Object temp = data.getInstances().get(0).get(nodeName);
        for (Instance instance : data.getInstances()) {
            if (!temp.equals(instance.get(nodeName))) {
                return null;
            }
        }
        return temp;
    }

    private void buildTree(Instances data, Object constributeValue, Node parentNode) {

        Attribute bestAttribute = getBestAttribute(data);

        Node newNode = new Node(bestAttribute, new HashMap<Object, Node>());
        parentNode.getBranchs().put(constributeValue, newNode);

        Object decideValue = tryGetDecideValue(data);
        if (decideValue != null) {
            // is leaf
            newNode.setDecideValue(decideValue);
            return;
        }

        Map<Object, Instances> childDataSets = buildChildDataSetsBaseOnBestAttribute(data,
                bestAttribute);
        for (Entry<Object, Instances> e : childDataSets.entrySet()) {
            buildTree(e.getValue(), e.getKey(), newNode);
        }
    }

    private Map<Object, Instances> buildChildDataSetsBaseOnBestAttribute(Instances data,
            Attribute bestAttribute) {

        // Build child dataSet
        //TODO do something to deal with real
        Map<Object, Instances> childDataSet = new HashMap<Object, Instances>();
        if (bestAttribute != null) {
            Map<Object, List<Instance>> instancesGroupedBySelectedAttributeValue = new HashMap<Object, List<Instance>>();
            for (Instance i : data.getInstances()) {
                List<Instance> instances = instancesGroupedBySelectedAttributeValue.get(i
                        .get(bestAttribute.getName()));
                if (instances == null) {
                    instances = new ArrayList<Instance>();
                    instancesGroupedBySelectedAttributeValue.put(i.get(bestAttribute.getName()),
                            instances);
                }
                instances.add(i);
            }

            for (Entry<Object, List<Instance>> e : instancesGroupedBySelectedAttributeValue
                    .entrySet()) {
                Instances instances = new Instances();
                instances.setInstances(e.getValue());

                List<Attribute> attrs = new ArrayList<Attribute>();
                for (Attribute attribute : data.getAttributes()) {
                    if (!attribute.getName().equals(bestAttribute.getName())) {
                        attrs.add(attribute);
                    }
                }
                instances.setAttributes(attrs);
                childDataSet.put(e.getKey(), instances);
            }
        }
        return childDataSet;
    }

    private Attribute getBestAttribute(Instances data) {
        // 1. Group cac thuoc tinh phan loai
        List<Object> classAttributeValues = data.getClassAttribute().getEnums();
        String classAttributeName = data.getClassAttribute().getName();

        // 2. Dem so luong mau cho tung group cua thuoc tinh phan loai
        // AttributeName/AttributeValue/ClassAttributeValue/Count
        Map<String, Map<Object, Map<Object, Integer>>> nbOfInstanceGroupedByAttributesAndClassAttribute = new HashMap<String, Map<Object, Map<Object, Integer>>>();
        for (Instance instance : data.getInstances()) {
            for (Attribute attr : data.getAttributes()) {

                if (!attr.equals(data.getClassAttribute())) {
                    //TODO: do something with this MAP to used in case of REAL
                    Map<Object, Map<Object, Integer>> attributeStatistic = nbOfInstanceGroupedByAttributesAndClassAttribute
                            .get(attr.getName());
                    if (attributeStatistic == null) {
                        attributeStatistic = new HashMap<Object, Map<Object, Integer>>();
                        nbOfInstanceGroupedByAttributesAndClassAttribute.put(attr.getName(),
                                attributeStatistic);
                    }

                    Map<Object, Integer> attributeValueStatistic = attributeStatistic.get(instance
                            .get(attr.getName()));
                    if (attributeValueStatistic == null) {
                        attributeValueStatistic = new HashMap<Object, Integer>();
                        attributeStatistic.put(instance.get(attr.getName()),
                                attributeValueStatistic);
                    }

                    for (Object classAttributeValue : classAttributeValues) {
                        if (attributeValueStatistic.get(classAttributeValue) == null) {
                            attributeValueStatistic.put(classAttributeValue, 0);
                        }
                        if (ObjectUtils.equals(instance.get(classAttributeName),
                                classAttributeValue)) {
                            Integer count = attributeValueStatistic.get(classAttributeValue);
                            count++;
                            attributeValueStatistic.put(instance.get(classAttributeName), count);
                            // break;
                        }
                    }
                }
            }
        }

        // 3. Tinh Gain Ratio cho tung thuoc tinh
        Map<String, Double> gainRatios = new HashMap<String, Double>();
        for (Attribute attr : data.getAttributes()) {
            if (!attr.equals(data.getClassAttribute())) {
                Double gain = getGainInfo(
                        nbOfInstanceGroupedByAttributesAndClassAttribute.get(attr.getName()), data);
                gainRatios.put(attr.getName(), gain);
            }
        }
        // 4. Chon thuoc tinh phan loai
        Attribute selectedAttribute = null;
        Double maxGainRaio = Double.MIN_VALUE;
        for (Attribute attr : data.getAttributes()) {
            if (!attr.equals(data.getClassAttribute())) {
                if (Double.compare(maxGainRaio, gainRatios.get(attr.getName())) < 0) {
                    maxGainRaio = gainRatios.get(attr.getName());
                    selectedAttribute = attr;
                }
            }
        }

        return selectedAttribute;
    }

    private Double getGainInfo(Map<Object, Map<Object, Integer>> attributeStatistic, Instances data) {
        Double nbOfInstance = data.getInstances().size() + 0d;
        Double result = getEntropyDataSet(data);
        for (Entry<Object, Map<Object, Integer>> e : attributeStatistic.entrySet()) {
            Integer count = 0;
            for (Entry<Object, Integer> sub_e : e.getValue().entrySet()) {
                count += sub_e.getValue();
            }
            Double temp = count / nbOfInstance;
            result -= temp * entropy(e.getValue());
        }
        return result;
    }

    private Double getEntropyDataSet(Instances data) {
        List<Object> classAttributeValues = data.getClassAttribute().getEnums();
        String classAttributeName = data.getClassAttribute().getName();
        Map<Object, Integer> statistics = new HashMap<Object, Integer>();
        for (Object classAttributeValue : classAttributeValues) {
            statistics.put(classAttributeValue, 0);
        }
        for (Instance instance : data.getInstances()) {
            statistics.put(instance.get(classAttributeName),
                    statistics.get(instance.get(classAttributeName)) + 1);
        }

        return entropy(statistics);
    }

    private Double entropy(Map<Object, Integer> value) {
        Integer count = 0;
        for (Entry<Object, Integer> e : value.entrySet()) {
            count += e.getValue();
        }
        Double result = 0d;
        for (Entry<Object, Integer> e : value.entrySet()) {
            Double temp = e.getValue() * 1.0 / count;
            if (Double.compare(temp, 0) != 0) {
                result -= temp * Math.log(temp) / log2;
            }
        }
        return result;
    }

    /**
     * {@inheritDoc}
     */
    public Object classifyInstance(Instance instance) {
        if (root == null) {
            throw new RuntimeException("Please train Decision tree before call classify");
        }
        Node node = root;
        while (true) {
            if (node == null) {
                return null;
            }
            if (node.getDecideValue() != null) {
                return node.getDecideValue();
            }
            if (node.getAttribute() == null) {
                return null;
            }
            node = node.getBranchs().get(instance.get(node.getAttribute().getName()));
        }
    }

    @Override
    public String toString() {
        return root.toString(0);
    }

    /**
     * @param args
     */
    public static void main(String[] args) {
        System.out.println("ID3 decision tree building...");
        long start = Calendar.getInstance().getTimeInMillis();
        System.out.println("Loading data set...");
        DataSource dataSource = new DataSource("D:/temp/tennis.csv.arff");
        Instances d = dataSource.getDataSet();
        System.out.println("Loaded data set: " + d.getRelation());
        ID3 classifier = new ID3();
        System.out.println("Building classifier: " + classifier.getClass());
        classifier.buildClassifer(d);
        long end = Calendar.getInstance().getTimeInMillis();
        System.out.println("The total training time: " + (end - start));
        System.out.println("--------------------------------------------------------------------");
        System.out.println("ID3 tree");
        System.out.println(classifier);
        System.out.println("--------------------------------------------------------------------");
        System.out.println("The training result:");
        for (Instance i : d.getInstances()) {
            System.out.print("ID: " + i.getId());
            System.out.print(", actual: "
                    + i.get(d.getClassAttribute().getName()));
            System.out.println(", predicted: " + classifier.classifyInstance(i));
        }
    }

}
